from typing import Sequence, Union

import numpy as np
import sapien.core as sapien
from sapien.core import Pose
from transforms3d.euler import euler2quat

from mani_skill2.agents.base_agent import BaseAgent
from mani_skill2.agents.configs.free_flying_robot import defaults


def get_entities_by_names(
    entities: Sequence[sapien.Entity], names: Union[str, Sequence[str]]
):
    assert isinstance(entities, (list, tuple)), type(entities)
    if isinstance(names, str):
        names = [names]
        flag = True  # singular
    else:
        flag = False
    ret = [None for _ in names]

    for entity in entities:
        name = entity.get_name()
        if name in names:
            ret[names.index(name)] = entity
    return ret[0] if flag else ret





class FreeFlyingRobot(BaseAgent):
    _config: defaults.FreeFlyingRobotDefaultConfig

    def __init__(
        self, scene, control_freq, control_mode=None, fix_root_link=True, config=None
    ):
        if control_mode is None:  # if user did not specify a control_mode
            control_mode = "pd_joint_delta_pos"
        super().__init__(
            scene,
            control_freq,
            control_mode=control_mode,
            fix_root_link=fix_root_link,
            config=config,
        )

    @classmethod
    def get_default_config(cls):
        return defaults.FreeFlyingRobotDefaultConfig()

    def _after_init(self):
        super()._after_init()

    #     (
    #         self.rfinger1_joint,
    #         self.rfinger2_joint,
    #         self.lfinger1_joint,
    #         self.lfinger2_joint,
    #     ) = get_entities_by_names(
    #         self.robot.get_joints(),
    #         [
    #             "right_panda_finger_joint1",
    #             "right_panda_finger_joint2",
    #             "left_panda_finger_joint1",
    #             "left_panda_finger_joint2",
    #         ],
    #     )
    #     (
    #         self.rfinger1_link,
    #         self.rfinger2_link,
    #         self.lfinger1_link,
    #         self.lfinger2_link,
    #     ) = get_entities_by_names(
    #         self.robot.get_links(),
    #         [
    #             "right_panda_leftfinger",
    #             "right_panda_rightfinger",
    #             "left_panda_leftfinger",
    #             "left_panda_rightfinger",
    #         ],
    #     )

    #     self.rhand, self.lhand = get_entities_by_names(
    #         self.robot.get_links(), ["right_panda_hand", "left_panda_hand"]
    #     )

    # def get_fingers_info(self):
    #     fingers_pos = self.get_ee_coords().flatten()
    #     fingers_vel = self.get_ee_vels().flatten()
    #     return {
    #         "fingers_pos": fingers_pos,
    #         "fingers_vel": fingers_vel,
    #     }

    # def get_ee_coords(self):
    #     finger_tips = [
    #         self.rfinger2_joint.get_global_pose().transform(Pose([0, 0.035, 0])).p,
    #         self.rfinger1_joint.get_global_pose().transform(Pose([0, -0.035, 0])).p,
    #         self.lfinger2_joint.get_global_pose().transform(Pose([0, 0.035, 0])).p,
    #         self.lfinger1_joint.get_global_pose().transform(Pose([0, -0.035, 0])).p,
    #     ]
    #     return np.array(finger_tips)

    # def get_ee_vels(self):
    #     finger_vels = [
    #         self.rfinger1_link.get_velocity(),
    #         self.rfinger2_link.get_velocity(),
    #         self.lfinger1_link.get_velocity(),
    #         self.lfinger2_link.get_velocity(),
    #     ]
    #     return np.array(finger_vels)
    # def get_ee_coords_sample(self, hand):
    #     if hand == "right":
    #         finger1_joint = self.rfinger1_joint
    #         finger2_joint = self.rfinger2_joint
    #     elif hand == "left":
    #         finger1_joint = self.lfinger1_joint
    #         finger2_joint = self.lfinger2_joint
    #     else:
    #         raise NotImplementedError
    #     l = 0.0355
    #     r = 0.052
    #     ret = []
    #     for i in range(10):
    #         x = (l * i + (4 - i) * r) / 4
    #         finger_tips = [
    #             finger2_joint.get_global_pose().transform(Pose([0, x, 0])).p,
    #             finger1_joint.get_global_pose().transform(Pose([0, -x, 0])).p,
    #         ]
    #         ret.append(finger_tips)
    #     return np.array(ret).transpose((1, 0, 2))

    # @staticmethod
    # def build_grasp_pose(approaching, closing, center):
    #     """Build a grasp pose (panda_hand)."""
    #     assert np.abs(1 - np.linalg.norm(approaching)) < 1e-3
    #     assert np.abs(1 - np.linalg.norm(closing)) < 1e-3
    #     assert np.abs(approaching @ closing) <= 1e-3
    #     ortho = np.cross(closing, approaching)
    #     T = np.eye(4)
    #     T[:3, :3] = np.stack([ortho, closing, approaching], axis=1)
    #     T[:3, 3] = center
    #     return Pose.from_transformation_matrix(T)

